from model.multimodel import model
from tools.infer import infer, infer_single
import torch
from loader.loader import get_loader
import colorama
from sklearn.metrics import roc_auc_score
import os
import csv


def list_files_in_directory(directory):
    try:
        files = os.listdir(directory)
        file_names = [file for file in files if os.path.isfile(os.path.join(directory, file))]
        return file_names
    except Exception as e:
        return str(e)


colorama.init(autoreset=True)
device = torch.device('cuda')
fold_index = 1
model.to(device)
model.eval()
# 定义文件夹路径
results_directory = 'results'
csv_file_path = os.path.join(results_directory, 'classification_5_results.csv')

# 检查文件夹是否存在，如果不存在则创建
if not os.path.exists(results_directory):
    os.makedirs(results_directory)

# 打开 CSV 文件以写入结果
with open(csv_file_path, mode='w', newline='') as csv_file:
    fieldnames = ['Fold', 'File', 'Accuracy', 'AUC']
    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
    writer.writeheader()  # 写入表头

    for fold_index in range(1, 6):
        directory = f'output5/fold_{fold_index}_classification/'
        dir_files = list_files_in_directory(directory)
        print(colorama.Fore.GREEN + f'fold {fold_index}')

        for file in dir_files:
            model.load_state_dict(torch.load(f'{directory}{file}'))
            _, _, test_loader = get_loader(False, fold_index)
            all_probs = []
            all_labels = []
            acc_num = total = 0

            with torch.no_grad():
                for images, labels in test_loader:
                    images, labels = images.to(device), labels.to(device)
                    outputs = model(images)
                    probs = torch.nn.functional.softmax(outputs, dim=1)
                    preds = probs.argmax(dim=1)
                    acc_num += torch.sum(preds == labels).cpu().item()
                    total += len(labels)
                    positive_probs = probs[:, 1]
                    all_probs.extend(positive_probs.cpu().numpy())  # 收集所有预测概率
                    all_labels.extend(labels.cpu().numpy())  # 收集所有真实标签

            # 计算准确率
            acc = acc_num / total * 100
            print(colorama.Fore.CYAN + f'{file} test acc: {acc:.2f}%', end='\t')

            # 计算AUC
            roc_auc = roc_auc_score(all_labels, all_probs)
            print(colorama.Fore.RED + f'{file} test AUC: {roc_auc:.4f}')

            # 将结果写入 CSV 文件
            writer.writerow({
                'Fold': fold_index,
                'File': file,
                'Accuracy': acc,
                'AUC': roc_auc
            })
